package com.example.demo.controller;

import com.example.demo.Mapper.DocumentMapper;
import com.example.demo.Mapper.SessionMapper;
import com.example.demo.context.BaseContext;
import com.example.demo.result.Result;
import com.example.demo.service.RagAssistant;
import com.example.demo.service.RagSerivce;
import com.example.demo.vo.documentDetailVO;
import com.example.demo.vo.messageDetailVO;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import lombok.RequiredArgsConstructor;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;

@RestController
@RequiredArgsConstructor
@RequestMapping("/api/rag")
@Tag(name="知识库接口")
public class RagController {

    final EmbeddingStore<TextSegment> embeddingStore;

    final EmbeddingModel embeddingModel;

    final DocumentMapper documentMapper;

    final RagSerivce ragSerivce;

    final SessionMapper sessionMapper;

//    @Operation(summary = "知识库ai会话接口")
//    @GetMapping("/high/chat")
//    public String chat(@RequestParam(value = "message")String message){
//        return assistant.chat(message);
//    }



//    @Operation(summary = "知识库ai会话接口")
//    @GetMapping("/chat")
//    public Result<messageDetailVO> highChat(@RequestParam(value = "userId")Long userId,
//                                            @RequestParam(value = "message")String message,
//                                            @RequestParam(required = false) Long sessionId){
//        return Result.success(ragSerivce.response(userId,message,sessionId));
//    }


    /**
     * 流式传输对话中的信息
     * @param message
     * @param sessionId
     * @return
     */
    @Operation(summary = "ai对话流式回复接口")
    @GetMapping(value = "/stream",produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<Object>> streamChat(@RequestParam(value = "message")String message,
                                                    @RequestParam(value = "sessionId") Long sessionId){
        return ragSerivce.streamResponse(message,sessionId);
    }

    @Operation(summary = "获取会话内所有信息接口")
    @GetMapping("/detail")
    public Result<List<messageDetailVO>> getMessageDetail(){
        List<messageDetailVO> messageDetailVOList=ragSerivce.getMessageDetail(BaseContext.getCurrentId());
        return Result.success(messageDetailVOList);
    }

    @Operation(summary = "获取会话id")
    @GetMapping("/getSessionId")
    public Result<Long> getSessionId(){
        Long sessionId= sessionMapper.getSessionId(BaseContext.getCurrentId());
        if(sessionId!=null){
            return Result.success(sessionId);
        }else {
            return Result.error("没有会话id");
        }
    }

    /**
     * 载入知识库文档
     * @return
     */
    @Operation(summary = "一键嵌入某用户所有知识库文档接口")
    @GetMapping("/load")
    public Result<String> load(){
        List<Document> documents=new ArrayList<>();
        System.out.println("用户id为"+BaseContext.getCurrentId());
        List<String> urlList=documentMapper.getUrl(BaseContext.getCurrentId());
        for(String url:urlList){
            Document document =UrlDocumentLoader.load(url,new ApacheTikaDocumentParser());
            //Document document = FileSystemDocumentLoader.loadDocument(url);
            documents.add(document);
        }
        //String docUrl="https://registration-system1.oss-cn-guangzhou.aliyuncs.com/txt.docx";
        //EmbeddingStoreIngestor.ingest(documents,embeddingStore);
        EmbeddingStoreIngestor.builder()
                .documentTransformer(document -> Document.document(document.text(),new Metadata().put("userId",BaseContext.getCurrentId())))
                //配置文本向量存储的位置
                .embeddingStore(embeddingStore)
                //配置文本向量模型
                .embeddingModel(embeddingModel)
                //怎么切分文本
                .documentSplitter(new DocumentByLineSplitter(100,30))
                .build().ingest(documents);//传入文本
        documentMapper.setEmbedding(BaseContext.getCurrentId());
        return Result.success("上传文档成功");
    }

    @Operation(summary = "根据知识库文档id进行嵌入接口")
    @GetMapping("/loadById")
    public Result<String> loadByDocumentId(@RequestParam("documentId") Long documentId){
        //判断这个文档id是否属于该用户
        com.example.demo.entity.Document document = ragSerivce.getById(documentId);
        String url=document.getUrl();
        Document document1 =UrlDocumentLoader.load(url,new ApacheTikaDocumentParser());
        EmbeddingStoreIngestor.builder()
                .documentTransformer(document2 -> Document.document(document2.text(),new Metadata().put("userId",BaseContext.getCurrentId())))
                //配置文本向量存储的位置
                .embeddingStore(embeddingStore)
                //配置文本向量模型
                .embeddingModel(embeddingModel)
                //怎么切分文本
                .documentSplitter(new DocumentByLineSplitter(50,30))
                .build().ingest(document1);//传入文本
        document.setIsEmbedding(1);
        documentMapper.updateById(document);
        return Result.success("上传文档成功");
    }

    /**
     * 获取某个用户的知识库文档
     * @param
     * @return
     */
    @Operation(summary = "获取知识库文档列表")
    @GetMapping("/list")
    public Result<List<documentDetailVO>> documentList(){
        return Result.success(ragSerivce.urlList(BaseContext.getCurrentId()));
    }

}
